from layer_norm import LayerNorm
import torch.nn as nn


class SublayerConnection(nn.Module):

    def __init__(self, feature_dim, dropout=0.1, norm_first=False):
        super(SublayerConnection, self).__init__()
        self.norm_first = norm_first
        self.norm = LayerNorm(feature_dim)
        self.dropout = nn.Dropout(p=dropout)

    def forward(self, x, sublayer):
        if self.norm_first:
            return x + self.dropout(sublayer(self.norm(x)))
        else:
            return x + self.dropout(self.norm(sublayer(x)))